package com.lndsnx.ai.controller;

import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.image.*;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URL;
import java.util.Collection;
import java.util.Set;
import java.util.stream.Collectors;


@RestController
@RequestMapping("/ai/image")
public class ImageModelController {

    @Resource
    private ImageModel imageModel;

    @GetMapping
    public void call(@RequestParam(value = "prompt", defaultValue = "为人工智能生成一张富有科技感的图片!", required = false) String prompt,
                     HttpServletResponse response) {
        ImageResponse imageResponse = imageModel.call(new ImagePrompt(prompt));
        String imageUrl = imageResponse.getResult().getOutput().getUrl();

        try {
            URL url = URI.create(imageUrl).toURL();
            InputStream in = url.openStream();

            response.setHeader("Content-Type", MediaType.IMAGE_PNG_VALUE);
            response.getOutputStream().write(in.readAllBytes());
            response.getOutputStream().flush();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @GetMapping("/multiPrompt")
    public ResponseEntity<Collection<String>> multiPrompt(@RequestParam(value = "prompt", defaultValue = "为人工智能生成一张富有科技感的图片!", required = false) String prompt,
                                                          @RequestParam(defaultValue = "2") int count,
                                                          HttpServletResponse response) {
        ImageOptions options = ImageOptionsBuilder.builder().N(count).build();
        ImageResponse imageResponse = imageModel.call(new ImagePrompt(prompt,options));
        Set<String> imageSet = imageResponse.getResults().stream().map(r -> r.getOutput().getUrl()).collect(Collectors.toSet());

        return ResponseEntity.ok(imageSet);
    }

}
